#include <rclcpp/rclcpp.hpp>
#include <geometry_msgs/msg/wrench_stamped.hpp>
#include <moveit/move_group_interface/move_group_interface.h>
#include <chrono>
#include <cmath>
#include <thread>

// ===== 导入 ONNX 力分类器 =====
#include "svm_predictor.hpp"

using namespace std::chrono_literals;

class ForceMoveLogger : public rclcpp::Node
{
public:
  ForceMoveLogger()
  : Node("force_move_logger"),
    latest_force_{0.0, 0.0, 0.0}
  {
    delay_after_motion_ms_ = 1300;  // 延迟时间(ms)

    // === 初始化分类器 ===
    classifier_ = std::make_unique<ForceClassifier>(
        "/home/jetson/zc6/src/svm/force_svm_model.onnx");

    RCLCPP_INFO(this->get_logger(), "✅ 分类器加载完成");

    // === 力传感器订阅 ===
    sub_wrench_ = this->create_subscription<geometry_msgs::msg::WrenchStamped>(
      "/wrench_comp", 10,
      std::bind(&ForceMoveLogger::wrenchCallback, this, std::placeholders::_1));

    RCLCPP_INFO(this->get_logger(), "节点初始化完成，2秒后开始运动...");
    timer_ = this->create_wall_timer(2s, std::bind(&ForceMoveLogger::startMotion, this));
  }

private:
  // ==================== 变量定义 ====================
  struct ForceValue {
    double fx, fy, fz;
  } latest_force_;

  int delay_after_motion_ms_;

  rclcpp::Subscription<geometry_msgs::msg::WrenchStamped>::SharedPtr sub_wrench_;
  rclcpp::TimerBase::SharedPtr timer_;
  std::unique_ptr<moveit::planning_interface::MoveGroupInterface> move_group_;
  std::unique_ptr<ForceClassifier> classifier_;

  // ==================== 启动运动 ====================
  void startMotion()
  {
    timer_->cancel();

    move_group_ = std::make_unique<moveit::planning_interface::MoveGroupInterface>(
        shared_from_this(), "zc4");
    move_group_->setMaxVelocityScalingFactor(0.3);
    move_group_->setMaxAccelerationScalingFactor(0.3);

    // Step 1: 回到 middle 位姿
    RCLCPP_INFO(this->get_logger(), "移动到 middle 位姿...");
    move_group_->setNamedTarget("middle");
    move_group_->move();

    // Step 2: 设定目标位姿
    geometry_msgs::msg::Pose target_pose;
    target_pose.position.x = 0.0362074;
    target_pose.position.y = 0.0;
    target_pose.position.z = 0.0708187;
    target_pose.orientation.x = -0.000396547;
    target_pose.orientation.y = 0.220008;
    target_pose.orientation.z = 0.0;
    target_pose.orientation.w = 0.975498;

    RCLCPP_INFO(this->get_logger(), "开始移动到目标位姿...");

    // ✅ 异步执行运动
    std::thread([this, target_pose]() {
      move_group_->setPoseTarget(target_pose);
      move_group_->move();  // 阻塞在子线程中
      RCLCPP_INFO(this->get_logger(), "机械臂运动完成，延时 %d ms ...",
                  delay_after_motion_ms_);
      std::this_thread::sleep_for(std::chrono::milliseconds(delay_after_motion_ms_));
      printInstantForce();
      rclcpp::shutdown();
    }).detach();
  }

  // ==================== 力回调 ====================
  void wrenchCallback(const geometry_msgs::msg::WrenchStamped::SharedPtr msg)
  {
    latest_force_.fx = msg->wrench.force.x;
    latest_force_.fy = msg->wrench.force.y;
    latest_force_.fz = msg->wrench.force.z;
  }

  // ==================== 力预测输出 ====================
  void printInstantForce()
  {
    double Fx = latest_force_.fx;
    double Fy = latest_force_.fy;
    double Fz = latest_force_.fz;

    // === 标准化参数（与 Python 训练一致） ===
    static const double mean_x = 1.72200128, mean_y = 1.67677145, mean_z = 1.71118453;
    static const double scale_x = 1.11652727, scale_y = 1.12525363, scale_z = 1.09691274;
    float Fx_n = (Fx - mean_x) / scale_x;
    float Fy_n = (Fy - mean_y) / scale_y;
    float Fz_n = (Fz - mean_z) / scale_z;

    // === 分类预测 ===
    int label = classifier_->predict(Fx_n, Fy_n, Fz_n);
    std::string category;
    if (label == 0) category = "hard";
    else if (label == 1) category = "medium";
    else if (label == 2) category = "soft";
    else category = "unknown";

    RCLCPP_INFO(this->get_logger(),
      "\n延时结束后瞬时力数据:\n"
      "Fx = %.4f N, Fy = %.4f N, Fz = %.4f N\n"
      "标准化后: (%.4f, %.4f, %.4f)\n"
      "分类结果: %s",
      Fx, Fy, Fz, Fx_n, Fy_n, Fz_n, category.c_str());
  }
};

// ==================== 主函数 ====================
int main(int argc, char **argv)
{
  rclcpp::init(argc, argv);

  // ✅ 使用多线程执行器以保证订阅与控制并行
  rclcpp::executors::MultiThreadedExecutor executor;
  auto node = std::make_shared<ForceMoveLogger>();
  executor.add_node(node);
  executor.spin();

  rclcpp::shutdown();
  return 0;
}
